﻿/*-
 * Copyright (C) 2018-2020, Shanghai Trusted Industry Control Platform Co., Ltd.
 * Created by 周正辉 (Zhenghui Zhou) <zhouzhenghui@ticpsh.com>, 赵楷文 (Kaiwen
 * Zhao) <zhaokaiwen@ticpsh.com>
 */

#include "testud.h"

#include "common/Coder.h"
#include "incparser/UDParser.h"
#include "testhelper.h"

using namespace std;
using namespace clang;
using namespace llvm;
using namespace rocket::common;
using njson = nlohmann::json;

std::string rocket::testif::TestUDTConsumer::parse_type(
    const clang::QualType &type_) {
  std::string type_name = getFullIFTypeName(type_, CI->getASTContext());
  clang::QualType type = type_;
  type.removeLocalRestrict();
  nlohmann::json::iterator type_it = types_j.find(type_name);
  if (type_it != types_j.end()) {
    assert(type_it.value().find("@canonical") != type_it.value().end());
    return type_it.value()["@canonical"];
  }

  std::string canonical_name = parse_canonical_type(type.getCanonicalType());
  if (type_name != canonical_name) {
    nlohmann::json attributes =
        types_j[canonical_name]["@attributes"];  // nlohmann::json::array()
    attributes.push_back("isNonCanonical");
    types_j[type_name]["@attributes"] = std::move(attributes);
    types_j[type_name]["@canonical"] = canonical_name;
    types_j[canonical_name]["@non-canonicals"].push_back(type_name);
    if (type.isLocalConstQualified()) {
      type.removeLocalConst();
      parse_type(type);
    }
    types_j[type_name]["@unqualified"] =
        types_j[canonical_name]["@unqualified"];
  }
  types_j[type_name]["@attributes"].push_back("isUserDefined");
  // elementCallBack(types[type_name]);
  return canonical_name;
}

static nlohmann::json createDiag(string type, string msg) {
  njson diag = njson::object();
  diag["type"] = type;
  diag["message"] = msg;
  return diag;
}
static nlohmann::json createDiag(
    string type,
    const rocket::incparser::UDParser::IllegalTypeParseException &e) {
  return createDiag(type, e.what());
}

void rocket::testif::TestUDTConsumer::handleEndSource() {
  // 构建原文件符号表
  TestTopLevelVisitor target(function_definitions, function_declarations,
                             global_declarations, CI->getASTContext(), setting,
                             (unsigned)setting.analyze_level);
  target.TraverseDecl(CI->getASTContext().getTranslationUnitDecl());

  // 检索模板函数
  struct UDTopLevelVisitor : TestTopLevelVisitor {
    virtual bool VisitFunctionTemplateDecl(
        clang::FunctionTemplateDecl *ftd) override {
      if (!ftd) return true;

      if (ftd && (ftd->getASTContext().getSourceManager().isInMainFile(
                      ftd->getLocation()) ||
                  analyze_level == AnalyzeLevel::AL_Full)) {
        addTargetFunction(ftd);
      }
      return true;
    }

    UDTopLevelVisitor(
        std::map<clang::FunctionDecl *, TargetFunction> &functions,
        std::map<const clang::FunctionDecl *, FuncDeclInfo> &functiondecls,
        std::map<clang::VarDecl *, TargetGlobalVar> &globals,
        clang::ASTContext &ast, const share::RocketAnalyzeSetting &setting,
        unsigned analyze_level)
        : TestTopLevelVisitor(functions, functiondecls, globals, ast, setting,
                              analyze_level) {}
  };

  // 模板函数原型检索
  class TemplFuncSearcher : public RecursiveASTVisitor<TemplFuncSearcher> {
   public:
    MangleContext *mc;
    std::map<std::string, FunctionTemplateDecl *> &specfs_name2decl;

    TemplFuncSearcher(
        ASTContext &ast,
        std::map<std::string, FunctionTemplateDecl *> &specfs_name2decl_)
        : mc(ast.createMangleContext()), specfs_name2decl(specfs_name2decl_) {}
    bool VisitFunctionTemplateDecl(FunctionTemplateDecl *ftd) {
      FunctionDecl *fd = ftd->getTemplatedDecl();
      std::string str = getFullNDeclName(fd, true);
      auto it = specfs_name2decl.find(str);
      if (it != specfs_name2decl.end()) {
        it->second = ftd;
      }
      return true;
    }

    bool VisitClassTemplateDecl(ClassTemplateDecl *ctd) {
      for (auto spec : ctd->specializations()) {
        for (auto spec_sub_decl : spec->decls()) {
          if (auto method = dyn_cast<FunctionDecl>(spec_sub_decl)) {
            this->VisitFunctionDecl(method);
          } else if (auto tmethod =
                         dyn_cast<FunctionTemplateDecl>(spec_sub_decl)) {
            this->VisitFunctionTemplateDecl(tmethod);
          }
        }
      }
      return true;
    }

    bool TraverseDecl(Decl *dl) {
      RecursiveASTVisitor<TemplFuncSearcher>::TraverseDecl(dl);
      return true;
    }
  };

  ///////////////////////////////////////////////////////////////////////

  njson &specs = udspecs_j;
  if (!specs.is_array()) {
    throw runtime_error(
        "the user define speclization actions information should be a list");
  }

  set<string> udspec_selected_ids;
  udspec_selected_ids.insert(udspec_selected_ids.begin(),
                             udspec_selected_ids.end());

  // 最后一次特化行为的id检查
  if (!specs.empty()) {
    auto last_it = specs.rbegin();
    auto id_it = last_it->find("@specid");
    if (id_it != last_it->end()) {
      udspec_selected_ids.insert(id_it->get<string>());
    } else {
      (*last_it)["@specid"] = "";
    }
  }

  /*解析增加的变量和函数信息*/
  udp.reset(new incparser::UDParser(CI));
  for (auto &spec_action_json : specs) {
    string current_spec_id;
    bool has_id = false;

    auto id_it = spec_action_json.find("@specid");
    if (id_it != spec_action_json.end()) {
      current_spec_id = id_it->get<string>();
      has_id = true;
    }

    if (has_id) setSpecID(current_spec_id);

    std::map<clang::FunctionDecl *, TargetFunction> inc_functions;
    std::map<const clang::FunctionDecl *, FuncDeclInfo> inc_functiondecls;
    std::map<clang::VarDecl *, TargetGlobalVar> inc_globals;

    // 增量解析
    UDTopLevelVisitor udvisitor(inc_functions, inc_functiondecls, inc_globals,
                                CI->getASTContext(), setting,
                                (unsigned)setting.analyze_level);

    if (spec_action_json.find("@type") != spec_action_json.end()) {
      // 用户定义类型特化
      QualType udt;
      try {
        if (spec_action_json.is_string()) {
          // - 字符串类型特化
          string udt_str = spec_action_json.get<string>();
          udt = udp->parseType(udt_str);
        } else if (spec_action_json.is_object()) {
          // 复杂结构特化
          auto type_j = spec_action_json.find("@type");
          auto args_j = spec_action_json.find("@templateArgs");
          if (args_j == spec_action_json.end()) {
            if (type_j != spec_action_json.end()) {
              // 无类模板实参直接增量编译类型
              string udt_str = type_j.value();
              udt = udp->parseType(udt_str);
            } else {
              throw runtime_error(
                  "failed to speclizate Template Arguements input json");
            }
          } else {
            auto targs = parseTemplateArgsJson(args_j.value());
            udt = udp->specCXXRecordTemplate(type_j.value().get<string>(),
                                             targs, true);
          }
        } else {
          throw runtime_error(
              "failed to speclizate Template Arguements input json");
        }
      } catch (incparser::UDParser::IllegalTypeParseException &e) {
        diags_j.push_back(createDiag(
            "error", "failed to specializate type with spec-action json:" +
                         spec_action_json.dump() +
                         ", with message:" + (string)e.what()));
      }

      udp->flushASTContext();

      if (udt->isRecordType()) {
        auto ud_record = udt->getAsRecordDecl();
        if (auto ud_cxxrd =
                dyn_cast<ClassTemplateSpecializationDecl>(ud_record)) {
          for (auto ud_cxxrd_method : ud_cxxrd->methods()) {
            if (!isa<CXXConstructorDecl>(ud_cxxrd_method)) {
              udvisitor.VisitFunctionDecl(ud_cxxrd_method);
            }
          }
          for (auto ftd : llvm::iterator_range<
                   DeclContext::specific_decl_iterator<FunctionTemplateDecl>>(
                   ud_cxxrd->decls())) {
            udvisitor.VisitFunctionTemplateDecl(ftd);
          }
        }

        // 解析内部特化情况
        udvisitor.TraverseDecl(ud_record);
      }

      parse_type(udt);

    } else if (spec_action_json.find("@mangled") !=
               spec_action_json.end()) {  // 整理函数特化信息
      std::map<std::string, FunctionTemplateDecl *>
          specfs_name2decl; /**< 模板函数原型*/
      std::map<std::string, std::vector<njson>>
          specfs_name2info; /**< 模板函数特化信息*/

      auto specf_ni_it = specfs_name2info.find(spec_action_json["@mangled"]);
      if (specf_ni_it == specfs_name2info.end()) {
        specfs_name2info[spec_action_json["@mangled"]] =
            std::vector<njson>{spec_action_json};
        specfs_name2decl[spec_action_json["@mangled"]] = nullptr;
      } else {
        specf_ni_it->second.push_back(spec_action_json);
      }

      {  // 检索原型
        TemplFuncSearcher tfs(CI->getASTContext(), specfs_name2decl);
        tfs.TraverseDecl(CI->getASTContext().getTranslationUnitDecl());
      }

      for (auto &info : specfs_name2info) {
        std::string mangled = info.first;
        auto &ele_arr = info.second;
        FunctionTemplateDecl *ftd;
        {
          auto dit = specfs_name2decl.find(mangled);
          if (dit == specfs_name2decl.end()) {
            throw runtime_error("could not specialize template function:" +
                                mangled + ", function declaration not found");
          }
          ftd = dit->second;
        }
        auto fd = ftd->getTemplatedDecl();

        // 增量编译特化函数指针
        for (auto ele_spec : ele_arr) {
          std::string name =
              "_ROCKET_INC_PARSE_FUNC_NAME_" + udp->getUniqueName();

          // 模板实参表
          vector<TemplateArgument> tas =
              parseTemplateArgsJson(ele_spec["@templateArgs"]);

          try {
            udp->specFunctionTemplate(fd, tas);
          } catch (incparser::UDParser::IllegalTypeParseException &e) {
            diags_j.push_back(
                createDiag("error",
                           "failed to specializate template function with "
                           "spec-action json:" +
                               spec_action_json.dump() +
                               ", with message:" + (string)e.what()));
          }
        }

        // 特化增量编译的模板实例
        udp->flushASTContext();
        udvisitor.TraverseDecl(ftd);
      }
    } else {
      throw runtime_error("Illegal user spec action, the json is " +
                          spec_action_json.dump());
    }

    /*
     *  新增内容的报文填充
     */

    // udp特化过程中产生的信息
    for (auto &udpmsg : udp->getRuntimeExceptions()) {
      diags_j.push_back(createDiag(
          "warning",
          "inside error message occured when specializate the template" +
              spec_action_json.dump() +
              ", with message:" + (string)udpmsg.what()));
    }
    udp->cleanExceptions();

    // 新增全局，填入键值
    for (auto &global : inc_globals) {
      std::string gname = getFullNDeclName(global.first);
      auto global_it = globals_j.find(gname);
      if (global_it == globals_j.end()) {
        auto j = parse_variable(global.first);
        auto &ja = j["@attributes"] = njson::array();
        ja.push_back("isUserDefined");
        if (has_id) j["@specid"] = current_spec_id;
        globals_j[gname] = std::move(j);
      }
    }

    // 新增函数，填入键值
    for (auto &inc_func : inc_functions) {
      auto incfunc_it = functions_j.begin();
      auto incfunc_it_end = functions_j.end();
      std::string incf_name = getFullNDeclName(inc_func.first, true);

      for (; incfunc_it != incfunc_it_end; incfunc_it++) {
        if ((*incfunc_it)["@mangled"] == incf_name) {
          break;
        }
      }

      if (incfunc_it != incfunc_it_end) {
        njson &incfunc_j = incfunc_it.value();
        vector<njson> incfunc_buf;
        for (auto &specf : inc_func.second.specfs) {
          if (this->function_declarations.count(specf.first->getFirstDecl())) {
            continue;
          }
          std::string spec_name = getFullNDeclName(specf.first, true);
          njson new_func = parse_function_definition(specf.first, specf.second);
          new_func["@attributes"].push_back("isUserDefined");
          incfunc_j["@templateSpecs"].push_back(spec_name);
          incfunc_buf.push_back(std::move(new_func));
        }
        for (auto &j : incfunc_buf) {
          functions_j.push_back(j);
        }
      } else {
        njson new_func =
            parse_function_definition(inc_func.first, inc_func.second);
        new_func["@attributes"].push_back("isUserDefined");
        functions_j.push_back(std::move(new_func));
      }
    }
  }

  /**
   * 过滤spec_id
   */
  if (udspec_selected_ids.count("all") == 0) {
    selectSpecId(types_j, udspec_selected_ids);
    selectSpecId(literals_j, udspec_selected_ids);
    selectSpecId(globals_j, udspec_selected_ids);
    selectSpecId(functions_j, udspec_selected_ids);
    //    selectSpecId(functiondecls_j, udspec_selected_ids);
  }
}

void rocket::testif::TestUDTConsumer::selectSpecId(
    nlohmann::json &j, const std::set<string> &id_set) {
  bool is_arr = j.is_array();

  for (auto e_it = j.begin(), e_end = j.end(); e_it != e_end;) {
    nlohmann::json &e_value = is_arr ? *e_it : e_it.value();
    bool changed = false;
    auto id_it = e_value.find("@specid");
    if (id_it != e_value.end()) {
      std::vector<std::string> change_it_list =
          id_it->get<std::vector<std::string>>();
      std::set<std::string> change_it_set;
      change_it_set.insert(change_it_list.begin(), change_it_list.end());

      for (auto &change_id : change_it_set) {
        if (change_id == "" || id_set.count(change_id)) {
          changed = true;
        }
      }
    }

    if (!changed) {
      e_it = j.erase(e_it);
      e_end = j.end();
    } else {
      ++e_it;
    }
  }
}

// TODO: 报文已被修改，重新实现
std::vector<TemplateArgument>
rocket::testif::TestUDTConsumer::parseTemplateArgsJson(
    const njson &targs_j) const {
  vector<TemplateArgument> tas;
  assert(targs_j.is_array());
  for (auto &targ : targs_j) {
    auto type_it = targ.find("@type");
    if (type_it != targ.end()) {
      auto expr_it = targ.find("@expr");
      if (expr_it != targ.end()) {
        // 非类型模板实参
        QualType ntta_t = udp->parseType(type_it->get<string>());

        std::string ntta_name =
            "__ROCKET_NON_TYPE_TEMPLATE_ARGUEMENT_NAME_" + udp->getUniqueName();
        auto init_expr = udp->parseInitlization(
            ntta_name, makeVarDecl(ntta_t, ntta_name, CI->getASTContext()) +
                           "=" + expr_it->get<string>() + ";");
        assert(init_expr);
        auto init_var_decl = dyn_cast<VarDecl>(init_expr);
        assert(init_var_decl);
        auto *init_decl = init_var_decl->getInit();
        assert(init_decl);

        if (IntegerLiteral *il = dyn_cast<IntegerLiteral>(init_decl)) {
          // 整数
          tas.push_back(TemplateArgument(il));
        } else {
          // 其他表达式
          tas.push_back(TemplateArgument(init_decl));
        }
      } else {
        // 类型模板实参
        QualType tta = udp->parseType(*type_it);
        tas.push_back(TemplateArgument(tta));
      }
    } else {
      // 模板模板实参
      auto name_it = targ.find("@template");
      assert(name_it != targ.end() && "not illegal template arguement");

      class TemplateRecordVisitor
          : public clang::RecursiveASTVisitor<TemplateRecordVisitor> {
        clang::ClassTemplateDecl *&ctd;
        const std::string &name;

       public:
        bool VisitClassTemplateDecl(ClassTemplateDecl *ctd_) {
          if (getFullNDeclName(ctd_->getTemplatedDecl()) == name) {
            ctd = ctd_;
          }
          return false;
        }
        TemplateRecordVisitor(clang::ClassTemplateDecl *&ctd_,
                              const std::string &name_)
            : ctd(ctd_), name(name_) {}
      };

      clang::ClassTemplateDecl *ctd = nullptr;
      {
        TemplateRecordVisitor trv(ctd, name_it->get<std::string>());
        trv.TraverseDecl(CI->getASTContext().getTranslationUnitDecl());
      }
      if (!ctd)
        throw runtime_error("Could not found the ClassTemplateDecl named " +
                            name_it->get<std::string>());
      tas.push_back(TemplateArgument(TemplateName(ctd)));
    }
  }
  return tas;
}
